﻿// petrk.cpp : 定义控制台应用程序的入口点。
//
#include "thisinc.h"
#include <string>
#include "wpetrk.h"
#include "murmur.h"
#include "minhasher.h"
#include "crc64.h"
#include "../../lib3rd/libjson/json.h"


enum { MIN_HASH_SIZE = 1000, LSH_BLOCKS = 200, BYTES_GRAM = 48 };

typedef wpe::Xmin_hasher<MIN_HASH_SIZE> min_hasher_t;

template < size_t GRAM >
long get_file_min_hash( wpe::ThunkFetcher & tfo, min_hasher_t & minhash, wpe::FileObject * fo )
{
	// 创建PE解析对象
	refp<wpe::PeParser> parser = new wpe::PeParser();
	if( !fo ) return E_OUTOFMEMORY;
	// 解析目标文件
	long lr = parser->open( fo );
	if( lr < 0 ) return E_FAIL;
	// 打开PE解析器
	lr = tfo.open( parser );
	if( lr < 0 ) return E_FAIL;
	// 开始提取
	uint64_t psh64 = 0;
	tfo.fetch();
	// 获得文本
	std::string tmp;
	tfo.get_as_text( tmp, false );
	//printf( "%s\n", tmp.c_str() );

	enum  { n_gram = GRAM };
	size_t strl = tmp.size();
	if( strl < n_gram ) return E_FAIL;

	const char * strp = tmp.c_str();

	minhash.reset();
	char gram[n_gram+1] = {};
	for( size_t i = 0; i < strl - n_gram; ++ i )
	{
		memcpy( gram, strp + i, n_gram );
		uint64_t hash64 = MurmurHash64B(gram,n_gram,0xdeadbeef);
		//printf( "%16I64X\n", hash64 );
		minhash.update( i, hash64 );
	}
	return S_OK;
}

template < size_t GRAM >
long get_file_min_hash( wpe::ThunkFetcher & tfo, min_hasher_t & minhash, const char * file )
{
	// 创建文件对象
	refp<wpe::FileObject> fo = new wpe::FileObject();
	if( !fo ) return E_OUTOFMEMORY;
	// 只读方式打开文件
	long lr = fo->open( file, "rb" );
	if( lr < 0 ) return E_FAIL;
	// 
	return get_file_min_hash<GRAM>( tfo, minhash, fo );
}

double compare_minhash( min_hasher_t & l, min_hasher_t & r )
{
	double same = 0;
	for( size_t i = 0; i < min_hasher_t::HASHES; ++ i )
	{
		if( l.hashes[i] == r.hashes[i] ) same ++;
	}
	return same / double(min_hasher_t::HASHES);
}

template < size_t NUM_IN_BLOCK >
long get_lsh_index_hashes( min_hasher_t & minhash, uint64_t * lsh, size_t hc )
{
	size_t c_minhash = __align_up(min_hasher_t::HASHES,NUM_IN_BLOCK);
	size_t c_blocks = c_minhash/NUM_IN_BLOCK;
	if( hc < c_blocks ) return E_INVALIDARG;
	size_t skip = NUM_IN_BLOCK * sizeof(uint64_t);
	const uint8_t * startp = (const uint8_t *)(minhash.hashes);
	const uint8_t * endp = (const uint8_t *)(minhash.hashes + min_hasher_t::HASHES);
	for( size_t i = 0; startp < endp; startp += skip, ++ i )
	{
		const uint8_t * this_endp = startp + skip;
		this_endp = MIN( this_endp, endp );
		lsh[i] = update_crc64( -1, startp, this_endp - startp );
	}
	return S_OK;
}

long compare_pe_files( const char * f1, const char * f2 )
{
	wpe::ThunkFetcher tfo1, tfo2;
	min_hasher_t h1, h2;

	long lr;

	lr = get_file_min_hash<BYTES_GRAM>( tfo1, h1, f1 );
	if( lr < 0 ) return lr;
	lr = get_file_min_hash<BYTES_GRAM>( tfo2, h2, f2 );
	if( lr < 0 ) return lr;

	for( size_t i = 0; i < min_hasher_t::HASHES; ++ i )
	{
		bool equl = (h1.hashes[i] == h2.hashes[i]);
		const char * eq = equl ? "==" : "!=";
#if TARGET_OS==OS_WIN
		printf( ":%04d | [%08X]%016I64X %s [%08X]%016I64X \n", i, h1.oindex[i], h1.hashes[i], eq, h2.oindex[i], h2.hashes[i] );
#else
		printf( ":%04d | %016llX %s %016llX \n", i, h1.hashes[i], eq, h2.hashes[i] );
#endif
	}

	printf( ":sim = %g\n", compare_minhash(h1,h2) );

	return 0;
}

long write_mag_file( const char * file, FILE * ofo )
{
	wpe::ThunkFetcher tfo;
	min_hasher_t mh;
	long lr = get_file_min_hash<BYTES_GRAM>( tfo, mh, file );
	if( lr < 0 ) return lr;
	uint64_t psh64;
	lr = tfo.get_psh( &psh64 );
	if( lr < 0 ) return lr;
	enum { LSH_HASH_PER_BLOCK = MIN_HASH_SIZE/LSH_BLOCKS };
	uint64_t lsh[LSH_BLOCKS] = {};
	lr = get_lsh_index_hashes<LSH_HASH_PER_BLOCK>(mh, lsh, LSH_BLOCKS );
	if( lr < 0 ) return lr;
	//////////////////////////////////////////////////////////////////////////
	char sz64[64] = {};
	Json::Value jval;
	sprintf( sz64, "%I64X", psh64 );
	jval["PSH"] = sz64;
	Json::Value& lshs = jval["LSH"];
	for( size_t i = 0; i < LSH_BLOCKS; ++ i )
	{
		Json::Value vlsh;
		sprintf( sz64, "%I64X", lsh[i] );
		vlsh = sz64;
		lshs.append( vlsh );
	}
	Json::Value& mhs = jval["MINHASH"];
	for( size_t i = 0; i < min_hasher_t::HASHES; ++ i )
	{
		Json::Value vmh;
		sprintf( sz64, "%I64X", mh.hashes[i] );
		vmh = sz64;
		mhs.append( vmh );
	}

	Json::StyledWriter w;
	std::string xxx = w.write( jval );
	printf( "%s\n", xxx.c_str() );
	fprintf( ofo, "%s\n", xxx.c_str() );

	return 0;
}

long write_mag_file( wpe::FileObject * fop, FILE * ofo )
{
	wpe::ThunkFetcher tfo;
	min_hasher_t mh;
	long lr = get_file_min_hash<BYTES_GRAM>( tfo, mh, fop );
	if( lr < 0 ) return lr;
	uint64_t psh64;
	lr = tfo.get_psh( &psh64 );
	if( lr < 0 ) return lr;
	enum { LSH_HASH_PER_BLOCK = MIN_HASH_SIZE/LSH_BLOCKS };
	uint64_t lsh[LSH_BLOCKS] = {};
	lr = get_lsh_index_hashes<LSH_HASH_PER_BLOCK>(mh, lsh, LSH_BLOCKS );
	if( lr < 0 ) return lr;
	//////////////////////////////////////////////////////////////////////////
	char sz64[64] = {};
	Json::Value jval;
	sprintf( sz64, "%I64X", psh64 );
	jval["PSH"] = sz64;
	Json::Value& lshs = jval["LSH"];
	for( size_t i = 0; i < LSH_BLOCKS; ++ i )
	{
		Json::Value vlsh;
		sprintf( sz64, "%I64X", lsh[i] );
		vlsh = sz64;
		lshs.append( vlsh );
	}
	Json::Value& mhs = jval["MINHASH"];
	for( size_t i = 0; i < min_hasher_t::HASHES; ++ i )
	{
		Json::Value vmh;
		sprintf( sz64, "%I64X", mh.hashes[i] );
		vmh.append( sz64 );
		vmh.append( mh.oindex[i] );
		mhs.append( vmh );
	}
	Json::StyledWriter w;
	std::string xxx = w.write( jval );
	//printf( "%s\n", xxx.c_str() );
	fprintf( ofo, "%s\n", xxx.c_str() );

	return 0;
}

/*

#pragma comment( lib, "Rpcrt4.lib" )

void gen_hash_seeds( size_t c )
{
	for( size_t i = 0; i < c; ++ i )
	{
		GUID guid;
		UuidCreate( &guid );
		uint64_t hash64 = MurmurHash64B( &guid, sizeof(guid), 0xdeadbeef);
		printf( "0x%016I64XLL,\n", hash64 );
	}
}

*/
//
//
//int main(int argc, char* argv[])
//{
//	//gen_hash_seeds( 20000 );
//	//return 0;
//
//	if( argc < 2 ) return E_INVALIDARG;
//
//	if( argc == 2 )
//	{
//		std::string file_name = argv[1];
//		file_name.append( ".trk" );
//		FILE * ofo = fopen( file_name.c_str(), "w" );
//		if( !ofo ) return E_FAIL;
//		long lr = write_mag_file( argv[1], ofo );
//		fclose( ofo );
//		return lr;
//	}
//	else if( argc == 3 )
//	{
//		return compare_pe_files( argv[1], argv[2] );
//	}
//	return E_INVALIDARG;
//}
//
//
//
